from typing import List, Dict, Any, Optional, Callable
from dataclasses import dataclass
import json
import re


@dataclass
class ContextSlot:
    """上下文槽位"""
    name: str
    content: str
    priority: int = 0
    metadata: Dict[str, Any] = None

    def __post_init__(self):
        if self.metadata is None:
            self.metadata = {}


class ContextOperator:
    """上下文操作符基类"""
    def process(self, slot: ContextSlot) -> ContextSlot:
        """处理槽位内容"""
        raise NotImplementedError


class InjectOperator(ContextOperator):
    """注入操作符"""

    def __init__(self, target_slot: str, content: str, position: str = "append"):
        self.target_slot = target_slot
        self.content = content
        self.position = position  # "prepend", "append", "replace"

    def process(self, slot: ContextSlot) -> ContextSlot:
        if slot.name == self.target_slot:
            if self.position == "prepend":
                new_content = self.content + "\n" + slot.content
            elif self.position == "append":
                new_content = slot.content + "\n" + self.content
            else:  # replace
                new_content = self.content

            return ContextSlot(
                name=slot.name,
                content=new_content,
                priority=slot.priority,
                metadata=slot.metadata
            )
        return slot


class InjectChainOperator(ContextOperator):
    """链式注入操作符"""

    def __init__(self, injections: List[Dict[str, Any]]):
        self.injections = injections
        self.inject_operators = []

        for injection in injections:
            operator = InjectOperator(
                target_slot=injection['target_slot'],
                content=injection['content'],
                position=injection.get('position', 'append')
            )
            self.inject_operators.append(operator)

    def process(self, slot: ContextSlot) -> ContextSlot:
        current_slot = slot
        for operator in self.inject_operators:
            current_slot = operator.process(current_slot)
        return current_slot


class CompressOperator(ContextOperator):
    """压缩操作符"""

    def __init__(self, compression_strategy: str = "smart", max_length: int = 2000):
        self.compression_strategy = compression_strategy
        self.max_length = max_length

    def process(self, slot: ContextSlot) -> ContextSlot:
        content = slot.content

        if self.compression_strategy == "truncate":
            compressed = self._truncate_compress(content)
        elif self.compression_strategy == "summary":
            compressed = self._summary_compress(content)
        elif self.compression_strategy == "smart":
            compressed = self._smart_compress(content)
        else:
            compressed = content

        return ContextSlot(
            name=slot.name,
            content=compressed,
            priority=slot.priority,
            metadata={**slot.metadata, "compressed": True}
        )

    def _truncate_compress(self, content: str) -> str:
        """简单截断压缩"""
        if len(content) <= self.max_length:
            return content

        # 保留开头和结尾的重要信息
        half_max = self.max_length // 2
        return content[:half_max] + "\n\n...[内容已压缩]...\n\n" + content[-half_max:]

    def _summary_compress(self, content: str) -> str:
        """摘要式压缩"""
        # 简单的基于规则的摘要生成
        sentences = re.split(r'[.!?。！？]+', content)
        important_sentences = []

        # 提取包含关键词的句子
        keywords = ['关键', '重要', '总结', '主要', '核心', '结论']
        for sentence in sentences:
            if any(keyword in sentence for keyword in keywords) and len(sentence.strip()) > 10:
                important_sentences.append(sentence.strip())

        if important_sentences:
            compressed = " | ".join(important_sentences[:5])
        else:
            # 如果没有找到关键词，取前几句和后几句
            compressed = self._truncate_compress(content)

        return compressed[:self.max_length]

    def _smart_compress(self, content: str) -> str:
        """智能压缩 - 结合多种策略"""
        if len(content) < self.max_length * 0.8:
            return content

        # 先尝试摘要压缩
        summary = self._summary_compress(content)
        if len(summary) < self.max_length:
            return summary

        # 如果摘要还是太长，使用截断
        return self._truncate_compress(content)


class ContextChain:
    """上下文链管理器"""

    def __init__(self):
        self.slots: Dict[str, ContextSlot] = {}
        self.operators: List[ContextOperator] = []
        self.execution_history: List[Dict] = []

    def add_slot(self, slot: ContextSlot):
        """添加上下文槽位"""
        self.slots[slot.name] = slot

    def add_operator(self, operator: ContextOperator):
        """添加操作符"""
        self.operators.append(operator)

    def inject(self, target_slot: str, content: str, position: str = "append"):
        """注入内容到指定槽位"""
        operator = InjectOperator(target_slot, content, position)
        self.add_operator(operator)
        self.execution_history.append({
            'type': 'inject',
            'target': target_slot,
            'position': position
        })

    def inject_chain(self, injections: List[Dict[str, Any]]):
        """链式注入多个内容"""
        operator = InjectChainOperator(injections)
        self.add_operator(operator)
        self.execution_history.append({
            'type': 'inject_chain',
            'injections': injections
        })

    def compress(self, compression_strategy: str = "smart", max_length: int = 2000):
        """添加压缩操作"""
        operator = CompressOperator(compression_strategy, max_length)
        self.add_operator(operator)
        self.execution_history.append({
            'type': 'compress',
            'strategy': compression_strategy,
            'max_length': max_length
        })

    def execute_chain(self) -> Dict[str, ContextSlot]:
        """执行整个操作链"""
        processed_slots = self.slots.copy()

        for operator in self.operators:
            for slot_name in list(processed_slots.keys()):
                processed_slots[slot_name] = operator.process(processed_slots[slot_name])

        return processed_slots

    def prepare_for_llm(self, slot_names: Optional[List[str]] = None) -> str:
        """准备最终注入到大模型的上下文"""
        processed_slots = self.execute_chain()

        if slot_names is None:
            slot_names = list(processed_slots.keys())

        # 按优先级排序
        sorted_slots = sorted(
            [processed_slots[name] for name in slot_names if name in processed_slots],
            key=lambda x: x.priority,
            reverse=True
        )

        # 构建最终上下文
        context_parts = []
        for slot in sorted_slots:
            if slot.content.strip():
                context_parts.append(f"【{slot.name}】\n{slot.content}")

        return "\n\n".join(context_parts)


# 使用示例
def demo_context_chain():
    """演示上下文链的完整流程"""

    # 1. 创建上下文链
    chain = ContextChain()

    # 2. 添加初始槽位
    chain.add_slot(ContextSlot(
        name="系统提示",
        content="你是一个有用的AI助手，请根据以下上下文回答问题。",
        priority=10
    ))

    chain.add_slot(ContextSlot(
        name="用户历史",
        content="用户之前询问过Python编程和机器学习相关问题。",
        priority=5
    ))

    chain.add_slot(ContextSlot(
        name="当前对话",
        content="",
        priority=8
    ))

    chain.add_slot(ContextSlot(
        name="知识库",
        content="这是一段很长的知识库内容..." * 100,  # 模拟长内容
        priority=3
    ))

    # 3. 执行操作序列
    # 注入当前用户问题
    chain.inject("当前对话", "用户提问：请解释MCP的工作原理。")

    # 链式注入多个相关内容
    chain.inject_chain([
        {
            'target_slot': '知识库',
            'content': 'MCP相关补充：Model Context Protocol是一种协议标准。',
            'position': 'append'
        },
        {
            'target_slot': '当前对话',
            'content': '补充上下文：MCP用于模型与外部工具交互。',
            'position': 'append'
        }
    ])

    # 压缩长内容
    chain.compress("smart", max_length=500)

    # 4. 准备最终上下文并注入LLM
    final_context = chain.prepare_for_llm()

    print("=== 最终上下文 ===")
    print(final_context)
    print(f"\n上下文长度: {len(final_context)} 字符")

    print("\n=== 执行历史 ===")
    for history in chain.execution_history:
        print(f"- {history}")

    return final_context


# 高级用法：自定义操作符
class CustomFilterOperator(ContextOperator):
    """自定义过滤操作符"""

    def __init__(self, filter_condition: Callable[[ContextSlot], bool]):
        self.filter_condition = filter_condition

    def process(self, slot: ContextSlot) -> ContextSlot:
        if not self.filter_condition(slot):
            # 如果不符合条件，清空内容但保留槽位
            return ContextSlot(
                name=slot.name,
                content="",
                priority=slot.priority,
                metadata=slot.metadata
            )
        return slot


if __name__ == "__main__":
    # 运行演示
    demo_context_chain()